package analisis;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import soot.Body;
import soot.Local;
import soot.SootMethod;
import soot.Unit;
import soot.UnitBox;
import soot.jimple.InvokeExpr;
import soot.jimple.InvokeStmt;
import soot.jimple.ParameterRef;
import soot.jimple.Stmt;
import soot.jimple.ThisRef;
import soot.jimple.internal.AbstractInstanceInvokeExpr;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIdentityStmt;
import soot.jimple.internal.JInvokeStmt;
import soot.jimple.internal.JReturnStmt;
import soot.jimple.internal.JimpleLocal;
import soot.toolkits.graph.BriefUnitGraph;
import soot.util.Chain;

public class SuperGraph extends BriefUnitGraph {

    public static Map<SootMethod, BriefUnitGraph> methodsCFGs;
    public static Map<Stmt, List<SootMethod>> methodsCalls;
    public static Unit mainMethodExitPoint = null;

    public SuperGraph(Body b) {
        super(b);
    }

    @Override
    protected void buildUnexceptionalEdges(Map<Unit, List<Unit>> unitToSuccs, Map<Unit, List<Unit>> unitToPreds) {

        mainMethodExitPoint = null;

        Map<SootMethod, Unit> methodsEntryPoints = new HashMap<SootMethod, Unit>();
        Map<SootMethod, Unit> methodsExitPoints = new HashMap<SootMethod, Unit>();

        Set<SootMethod> methods = methodsCFGs.keySet();

        //guardamos entries y exits de cada método para saber dónde conectar después
        for (SootMethod m : methods) {
            BriefUnitGraph methodUnitGraph = methodsCFGs.get(m);
            Iterator<Unit> entryExitIt = methodUnitGraph.iterator();
            Unit entryExit = entryExitIt.next();
            methodsEntryPoints.put(m, entryExit);
            while (entryExitIt.hasNext()) {
                entryExit = entryExitIt.next();
            }
            methodsExitPoints.put(m, entryExit);
        }

        mainMethodExitPoint = methodsExitPoints.get(this.body.getMethod());

        //Se agregan todas las Units de todos los metodos (menos el actual)
        //al final del unit chain del metodo actual
        for (SootMethod m : methods) {
            BriefUnitGraph methodUnitGraph = methodsCFGs.get(m);

            if (m != this.body.getMethod()) {
                for (Unit u : methodUnitGraph) {
                    unitChain.add(u);
                }
            }
        }

        //renombramos todas las variables para que tengan nombre del método antes
        for (SootMethod m : methods) {
            for (Local l : m.getActiveBody().getLocals()) {
                l.setName(
                        ((l.getName().startsWith("$")) ? "$" : "") +
                        m.getName() + "_" + l.getName());
            }
        }

        //bindeo de parámetros y return value
        for (Iterator<Unit> ui = unitChain.snapshotIterator(); ui.hasNext();) {
            Unit u = ui.next();
            List<SootMethod> methodsToCall = methodsCalls.get((Stmt)u);
            if (methodsToCall != null) {
                for (SootMethod m : methodsToCall) {
                    InvokeExpr invokeExpr;

                    //si el método es void, el nodo es un JInvokeStmt directo
                    if (u instanceof JInvokeStmt) {
                        invokeExpr = ((JInvokeStmt) u).getInvokeExpr();
                    } //si devuelve algo es un JAssignStmt y la expr está a la derecha
                    else {
                        invokeExpr = (InvokeExpr) ((JAssignStmt) u).getRightOp();
                    }

                    //bindeo de this (parámetro implícito)

                    //nombre de this en la llamada
                    JimpleLocal thisLocal = (JimpleLocal) ((AbstractInstanceInvokeExpr) invokeExpr).getBase();
                    //nombre de this adentro del método
                    JimpleLocal thisRemote = (JimpleLocal) ((JIdentityStmt) m.getActiveBody().getUnits().iterator().next()).getLeftOp();

                    //agrego el nodo thisRemote = thisLocal
                    JAssignStmt assignStmtThis = new JAssignStmt(thisRemote, thisLocal);
                    unitChain.insertBefore(assignStmtThis, u);

                    //si tiene parámetros (explícitos)
                    if (m.getParameterCount() > 0) {
                        BriefUnitGraph methodCFG = methodsCFGs.get(m);

                        int i = 0;
                        for (Unit possibleParameter : methodCFG) {
                            if (possibleParameter instanceof JIdentityStmt &&
                                    ((JIdentityStmt) possibleParameter).getRightOp() instanceof ParameterRef) {

                                //este el parametro que se le pasa
                                JimpleLocal argument = (JimpleLocal) invokeExpr.getArg(i);

                                //este es el objeto del parámetro en el método invocado
                                JimpleLocal parameterName = (JimpleLocal) ((JIdentityStmt) possibleParameter).getLeftOp();

                                //bindeo de parámetro: parameterName = argument
                                JAssignStmt assignStmt = new JAssignStmt(parameterName, argument);
                                //inserto nodo ANTES del call
                                unitChain.insertBefore(assignStmt, u);
                                i++;
                            }
                            //si encuentro un return y no es void
                            if (possibleParameter instanceof JReturnStmt && u instanceof JAssignStmt) {
                                JReturnStmt returnStmt = (JReturnStmt) possibleParameter;
                                //qué devuelve
                                JimpleLocal returnedVar = (JimpleLocal) returnStmt.getOp();
                                //adónde lo asigna
                                JimpleLocal assignedTo = (JimpleLocal) ((JAssignStmt) u).getLeftOp();

                                //bindeo return: assignedTo = returnedVar
                                JAssignStmt assignStmt = new JAssignStmt(assignedTo, returnedVar);
                                //inserto nodo DESPUÉS del call
                                unitChain.insertAfter(assignStmt, u);
                            }
                        }
                    }
                }

            }
        }

        //sacamos todos los @parameterX y @this, para que no pisen los que seteamos
        //también sacamos llamadas a <init> (constructores)
        for (Iterator<Unit> ui = unitChain.snapshotIterator(); ui.hasNext();) {
            Unit u = ui.next();
            if (u instanceof JIdentityStmt) {
                JIdentityStmt idStmt = (JIdentityStmt) u;
                if (idStmt.getRightOp() instanceof ParameterRef) {
                    unitChain.remove(u);
                }

                if (idStmt.getRightOp() instanceof ThisRef) {
                    //acá no puedo sacar el nodo porque es el entrypoint, le cambio el valor
                    //de la izquierda a TRAP para que no me pise el pasaje implícito
                    idStmt.setLeftOp(new JimpleLocal("TRAP", soot.IntType.v()));
                }
            }
            if (u instanceof InvokeStmt) {
                if (((InvokeStmt) u).getInvokeExpr().getMethod().getName().equals("<init>")) {
                    unitChain.remove(u);
                }
            }
        }

        /**
         * <código original>
         */
        // Initialize the predecessor sets to empty
        for (Iterator<Unit> unitIt = unitChain.iterator(); unitIt.hasNext();) {
            unitToPreds.put(unitIt.next(), new ArrayList<Unit>());
        }

        Iterator<Unit> unitIt = unitChain.iterator();
        Unit currentUnit, nextUnit;

        nextUnit = unitIt.hasNext() ? (Unit) unitIt.next() : null;

        while (nextUnit != null) {
            currentUnit = nextUnit;
            nextUnit = unitIt.hasNext() ? (Unit) unitIt.next() : null;

            List<Unit> successors = new ArrayList<Unit>();

            if (currentUnit.fallsThrough()) {
                // Add the next unit as the successor
                if (nextUnit != null) {
                    successors.add(nextUnit);
                    unitToPreds.get(nextUnit).add(currentUnit);
                }
            }

            if (currentUnit.branches()) {
                for (Iterator<UnitBox> targetIt = currentUnit.getUnitBoxes().iterator(); targetIt.hasNext();) {
                    Unit target = targetIt.next().getUnit();
                    // Arbitrary bytecode can branch to the same
                    // target it falls through to, so we screen for duplicates:
                    if (!successors.contains(target)) {
                        successors.add(target);
                        unitToPreds.get(target).add(currentUnit);
                    }
                }
            }

            // Store away successors
            unitToSuccs.put(currentUnit, successors);
        }

        /**
         * </código original>
         */
        //limpiamos succ y pred de llamadas a métodos
        for (Iterator<Unit> unitIt2 = unitChain.iterator(); unitIt2.hasNext();) {

            Unit currUn = unitIt2.next();
            Unit nextUn = getNextUnit(unitChain, currUn);

            List<SootMethod> methodsToCall = methodsCalls.get(currUn);
            if (methodsToCall != null) {
                unitToSuccs.get(currUn).clear();
                unitToPreds.get(nextUn).clear();
            }
        }

        //conectamos para las llamadas
        for (Iterator<Unit> unitIt2 = unitChain.iterator(); unitIt2.hasNext();) {

            Unit currUn = unitIt2.next();
            Unit nextUn = getNextUnit(unitChain, currUn);

            List<SootMethod> methodsToCall = methodsCalls.get((Stmt)currUn);
            if (methodsToCall != null) {
                for (SootMethod methodToCall : methodsToCall) {
                    Unit methodEntry = methodsEntryPoints.get(methodToCall);
                    Unit methodExit = methodsExitPoints.get(methodToCall);

                    //conectamos llamada con entry del método
                    unitToSuccs.get(currUn).add(methodEntry);
                    unitToPreds.get(methodEntry).add(currUn);

                    //conectamos exit del método con siguiente unit
                    unitToSuccs.get(methodExit).add(nextUn);
                    unitToPreds.get(nextUn).add(methodExit);
                }
            }
        }

        //dump de cómo quedó
        for (Iterator<Unit> ui = unitChain.snapshotIterator(); ui.hasNext();) {
            Unit u = ui.next();
            for (SootMethod m : methodsEntryPoints.keySet()) {
                if (methodsEntryPoints.get(m) == u) {
                    System.out.println(m.getName());
                    break;
                }
            }

            String preds = "", succs = "";
            for (Unit uu : getPredsOf(u)) {
                preds += uu.hashCode() + ", ";
            }
            for (Unit uu : getSuccsOf(u)) {
                succs += uu.hashCode() + ", ";
            }

            //System.out.println("       pred: "  + preds);
            System.out.println("   " + u.hashCode() + ": " + u);
            //System.out.println("       succ: "  + succs);
        }
    }

    //devuelve el Unit siguiente en la lista, null si no hay siguiente
    protected Unit getNextUnit(Chain<Unit> unitChain, Unit u) {
        for (Iterator<Unit> u2it = unitChain.iterator(); u2it.hasNext();) {
            Unit u2 = u2it.next();

            if (!u2it.hasNext()) {
                return null;
            }
            if (u.equals(u2)) {
                return u2it.next();
            }
        }
        return null;
    }
}
